#include <arpa/inet.h>
#include <errno.h>
#include <fcntl.h>
#include <netinet/in.h>
#include <pthread.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/epoll.h>
#include <sys/socket.h>
#include <unistd.h>

#define MAX_EVENTS   1024
#define BUFFER_SIZE  4096
#define THREAD_COUNT 4  // 线程数量，应与CPU核心数匹配

// 线程参数结构
typedef struct {
    int thread_id;
    const char* ip;
    int port;
} thread_param_t;

// 设置套接字为非阻塞模式
int set_nonblocking(int fd)
{
    int flags = fcntl(fd, F_GETFL, 0);
    if (flags == -1)
        return -1;
    return fcntl(fd, F_SETFL, flags | O_NONBLOCK);
}

// 创建监听套接字并设置SO_REUSEPORT
int create_listen_socket(const char* ip, int port)
{
    int listen_fd;
    struct sockaddr_in server_addr;
    int optval = 1;

    // 创建套接字
    if ((listen_fd = socket(AF_INET, SOCK_STREAM, 0)) < 0) {
        perror("socket creation failed");
        return -1;
    }

    // 设置SO_REUSEADDR和SO_REUSEPORT
    if (setsockopt(listen_fd, SOL_SOCKET, SO_REUSEADDR, &optval, sizeof(optval)) < 0) {
        perror("setsockopt SO_REUSEADDR failed");
        close(listen_fd);
        return -1;
    }

    if (setsockopt(listen_fd, SOL_SOCKET, SO_REUSEPORT, &optval, sizeof(optval)) < 0) {
        perror("setsockopt SO_REUSEPORT failed");
        close(listen_fd);
        return -1;
    }

    // 绑定地址
    memset(&server_addr, 0, sizeof(server_addr));
    server_addr.sin_family = AF_INET;
    server_addr.sin_addr.s_addr = inet_addr(ip);
    server_addr.sin_port = htons(port);

    if (bind(listen_fd, (struct sockaddr*)&server_addr, sizeof(server_addr)) < 0) {
        perror("bind failed");
        close(listen_fd);
        return -1;
    }

    // 开始监听
    if (listen(listen_fd, SOMAXCONN) < 0) {
        perror("listen failed");
        close(listen_fd);
        return -1;
    }

    // 设置为非阻塞模式
    if (set_nonblocking(listen_fd) < 0) {
        perror("set_nonblocking failed");
        close(listen_fd);
        return -1;
    }

    return listen_fd;
}

// 处理客户端连接
void handle_client(int client_fd)
{
    char buffer[BUFFER_SIZE];
    ssize_t bytes_read;

    // 读取数据
    bytes_read = read(client_fd, buffer, BUFFER_SIZE - 1);
    if (bytes_read > 0) {
        buffer[bytes_read] = '\0';
        printf("Received: %s\n", buffer);

        // 简单回显响应
        write(client_fd, buffer, bytes_read);
    }

    close(client_fd);
}

// 工作线程函数
void* worker_thread(void* arg)
{
    thread_param_t* param = (thread_param_t*)arg;
    int thread_id = param->thread_id;
    const char* ip = param->ip;
    int port = param->port;

    // 每个线程创建自己的监听套接字
    int listen_fd = create_listen_socket(ip, port);
    if (listen_fd < 0) {
        fprintf(stderr, "Thread %d: Failed to create listen socket\n", thread_id);
        pthread_exit(NULL);
    }

    // 创建epoll实例
    int epoll_fd = epoll_create1(0);
    if (epoll_fd == -1) {
        perror("epoll_create1");
        close(listen_fd);
        pthread_exit(NULL);
    }

    // 将监听套接字添加到epoll
    struct epoll_event ev;
    ev.events = EPOLLIN | EPOLLET;
    ev.data.fd = listen_fd;

    if (epoll_ctl(epoll_fd, EPOLL_CTL_ADD, listen_fd, &ev) == -1) {
        perror("epoll_ctl: listen_fd");
        close(listen_fd);
        close(epoll_fd);
        pthread_exit(NULL);
    }

    printf("Thread %d started with listen FD %d and epoll FD %d\n", thread_id, listen_fd, epoll_fd);

    struct epoll_event events[MAX_EVENTS];
    int nfds, i;
    int client_fd;
    struct sockaddr_in client_addr;
    socklen_t client_len = sizeof(client_addr);

    while (1) {
        nfds = epoll_wait(epoll_fd, events, MAX_EVENTS, -1);
        if (nfds == -1) {
            perror("epoll_wait");
            continue;
        }

        for (i = 0; i < nfds; i++) {
            // 处理新连接
            if (events[i].data.fd == listen_fd) {
                while ((client_fd = accept(listen_fd, (struct sockaddr*)&client_addr, &client_len))
                       > 0) {
                    printf("Thread %d accepted connection from %s:%d\n", thread_id,
                           inet_ntoa(client_addr.sin_addr), ntohs(client_addr.sin_port));

                    // 设置为非阻塞模式并添加到epoll
                    set_nonblocking(client_fd);

                    struct epoll_event client_ev;
                    client_ev.events = EPOLLIN | EPOLLET;
                    client_ev.data.fd = client_fd;

                    if (epoll_ctl(epoll_fd, EPOLL_CTL_ADD, client_fd, &client_ev) == -1) {
                        perror("epoll_ctl: client_fd");
                        close(client_fd);
                    }
                }

                if (errno != EAGAIN && errno != EWOULDBLOCK) {
                    perror("accept");
                }
            } else {
                // 处理客户端数据
                handle_client(events[i].data.fd);
            }
        }
    }

    close(listen_fd);
    close(epoll_fd);
    return NULL;
}

int main(int argc, char* argv[])
{
    if (argc != 3) {
        fprintf(stderr, "Usage: %s <IP> <PORT>\n", argv[0]);
        exit(EXIT_FAILURE);
    }

    const char* ip = argv[1];
    int port = atoi(argv[2]);
    pthread_t threads[THREAD_COUNT];
    thread_param_t thread_params[THREAD_COUNT];
    int i;

    // 创建线程
    for (i = 0; i < THREAD_COUNT; i++) {
        // 设置线程参数
        thread_params[i].thread_id = i;
        thread_params[i].ip = ip;
        thread_params[i].port = port;

        // 创建线程
        if (pthread_create(&threads[i], NULL, worker_thread, &thread_params[i]) != 0) {
            perror("pthread_create");
            exit(EXIT_FAILURE);
        }
    }

    printf("Server started on %s:%d with %d threads\n", ip, port, THREAD_COUNT);

    // 等待所有线程结束
    for (i = 0; i < THREAD_COUNT; i++) {
        pthread_join(threads[i], NULL);
    }

    return 0;
}
